from keras.utils import np_utils
from sklearn.model_selection import train_test_split
import numpy as np
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import Flatten
from keras.layers.convolutional import Conv2D
from keras.layers.convolutional import MaxPooling2D
from matplotlib import pyplot as plt
from PIL import Image
# 加载数据集
data = np.load("数据集X.npy")
label = np.load("数据集Y.npy")

# 将label标签进行one-hot编码
label = np_utils.to_categorical(label)

# 划分训练集合和测试集
X_train, X_test, Y_train, Y_test = train_test_split(data, label, test_size=0.2)
# 构建CNN模型
model = Sequential()
model.add(Conv2D(16, (3, 3), activation='relu', input_shape=(28, 28, 1)))
model.add(MaxPooling2D((2, 2)))
model.add(Conv2D(32, (5, 5), activation='relu'))
model.add(MaxPooling2D((2, 2)))
model.add(Flatten())
model.add(Dropout(0.5))
model.add(Dense(128, activation='relu'))
model.add(Dense(10, activation='softmax'))

# 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
# 训练模型
model.fit(X_train, Y_train, epochs=100,  batch_size=5, validation_data=(X_test, Y_test))

# 进行预测
img = Image.open("data/测试图片2.png").convert("L").resize((28, 28))
plt.imshow(img)
plt.show()
img = np.array(img).reshape(1, 28, 28, 1)

img = img.astype("float32")/255.0
print('结果:', np.argmax(model.predict(img)))
